import rclpy
from rclpy.node import Node
import threading
from topic_msg.msg import CatVision
import cv2
import face_recognition
import numpy as np
import time


class CatVisionNode(Node):
    def __init__(self, name):
        super().__init__(name)
        self._upsample_times = 1
        self._model = "hog"
        self._publisher = self.create_publisher(CatVision, "cat_vision", 10)
        # pipeline = "v4l2src device=/dev/video0 ! videoconvert ! appsink"
        self._cap = cv2.VideoCapture(0)
        if not self._cap.isOpened():
            self.get_logger().info("管道初始化失败！请检查语法或依赖。")
            exit()
        self._thd = threading.Thread(target=self._run)
        self._thd.start()

    def __del__(self):
        self._thd.join()
        self._cap.release()
        self._publisher.destroy()

    def _run(self):
        while rclpy.ok():
            ret, img = self._cap.read()
            if ret is None:
                break
            self._deal_image(img)
            time.sleep(0.1)
            # key = cv2.waitKey(200)
            # if key == 220:
            # break

    def _deal_image(self, image):
        msg = self._face_reg(image)
        if msg.number != 0:
            msg.type = CatVision.FACE_TYPE
        else:
            # msg = self._ball_reg(image)
            # if msg.number != 0:
            #     msg.type = CatVision.BALL_TYPE
            # else:
            #     msg.type = CatVision.EMPTY_TYPE
            msg.type = CatVision.EMPTY_TYPE
        height, width = image.shape[:2]
        msg.width = int(width)
        msg.height = int(height)
        # self.get_logger().info(f"image hight/width {msg.height }/{msg.width}")
        self._publisher.publish(msg)

    def _face_reg(self, image):
        "摄像头上圆圈标记表示视野原点坐标，安装时将其置于左上方"
        msg = CatVision()
        face_locations = face_recognition.face_locations(
            image, self._upsample_times, self._model
        )
        msg.number = len(face_locations)
        if msg.number == 0:
            return msg
        msg.type = CatVision.FACE_TYPE
        for (
            top,
            right,
            bottom,
            left,
        ) in face_locations:
            msg.y.append(int((top + bottom) / 2))
            msg.x.append(int((right + left) / 2))
        return msg

    def _ball_reg(self, image):
        msg = CatVision()
        # 1. 转换为灰度图并进行模糊降噪
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        gray = cv2.medianBlur(gray, 5)  # 中值滤波有效去除椒盐噪声

        # 2. 霍夫圆检测
        circles = cv2.HoughCircles(
            image=gray,
            method=cv2.HOUGH_GRADIENT,  # 检测方法（目前仅支持梯度法）
            dp=1,  # 累加器分辨率（1=与原图相同）
            minDist=20,  # 圆之间的最小距离
            param1=50,  # Canny边缘检测高阈值
            param2=30,  # 累加器阈值（越小检测越多假圆）
            minRadius=10,  # 最小圆半径
            maxRadius=100,  # 最大圆半径
        )

        # 4. 处理检测结果
        if circles is None:
            msg.number = 0
            return msg
        msg.type = CatVision.BALL_TYPE
        msg.number = 0
        circles = np.round(circles[0, :]).astype("int")
        for x, y, r in circles:
            if r > 90:
                msg.x.append(x)
                msg.y.append(y)
                msg.number += 1
        return msg


def main():
    rclpy.init()
    node = CatVisionNode("cat_vision")
    rclpy.spin(node)
    rclpy.shutdown()
